import logging
import json
from typing import List, Dict, Any, Union, Optional
from langchain_community.llms import Ollama
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser

class ChartGenerator:
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.llm = Ollama(model="qwen2.5:latest")
        self.supported_chart_types = [
            "bar", "line", "pie"
        ]
    
    def generate_chart_config(self, data: List[Dict[str, Any]], user_query: str, chart_type: Optional[str] = None) -> Dict[str, Any]:
        try:
            if not chart_type:
                chart_type = self._recommend_chart_type(data, user_query)
            if chart_type == "bar":
                return self._generate_bar_chart(data, user_query)
            elif chart_type == "line":
                return self._generate_line_chart(data, user_query)
            elif chart_type == "pie":
                return self._generate_pie_chart(data, user_query)
        except Exception as e:
            self.logger.error(f"Error generating chart config: {str(e)}")
    
    def _recommend_chart_type(self, data: List[Dict[str, Any]], user_query: str) -> str:
        data_features = self._analyze_data_features(data)
        prompt = PromptTemplate.from_template(
            """作为数据可视化专家，您需要为给定数据推荐最合适的图表类型。
            数据特征：
            {data_features}
            
            用户查询：{user_query}
            
            请从以下选项中选择最合适的图表类型：
            - bar: 柱状图，适用于类别比较
            - line: 折线图，适用于时间趋势或连续变化
            - pie: 饼图，适用于显示比例或构成
            
            只返回英文图表类型名称，不要提供任何解释。
            """
        )
        
        chain = prompt | self.llm | StrOutputParser()
        chart_type = chain.invoke({
            "data_features": json.dumps(data_features, ensure_ascii=False, indent=2),
            "user_query": user_query
        }).strip().lower()
        return chart_type
    
    def _analyze_data_features(self, data: List[Dict[str, Any]]) -> Dict[str, Any]:
        if not data:
            return {"empty": True}
        columns = list(data[0].keys())
        column_types = {}
        numeric_columns = []
        categorical_columns = []
        for col in columns:
            sample_values = [row[col] for row in data if col in row]
            if not sample_values:
                continue
            is_numeric = all(
                isinstance(val, (int, float)) or 
                (isinstance(val, str) and val.replace('.', '', 1).isdigit()) 
                for val in sample_values if val is not None
            )
            if is_numeric:
                column_types[col] = "numeric"
                numeric_columns.append(col)
            else:
                column_types[col] = "categorical"
                categorical_columns.append(col)
        column_count = len(columns)
        has_category_value_pairs = len(categorical_columns) >= 1 and len(numeric_columns) >= 1
        return {
            "column_count": column_count,
            "columns": columns,
            "column_types": column_types,
            "numeric_columns": numeric_columns,
            "categorical_columns": categorical_columns,
            "has_category_value_pairs": has_category_value_pairs
        }
    
    def _generate_bar_chart(self, data: List[Dict[str, Any]], user_query: str) -> Dict[str, Any]:
        features = self._analyze_data_features(data)
        x_column = features["categorical_columns"][0] if features["categorical_columns"] else features["columns"][0]
        y_columns = features["numeric_columns"] if features["numeric_columns"] else [features["columns"][1]] if len(features["columns"]) > 1 else [features["columns"][0]] 
        x_data = [str(row[x_column]) for row in data if x_column in row]
        series = []
        for y_column in y_columns:
            series.append({
                "name": y_column,
                "type": "bar",
                "data": [row.get(y_column, 0) for row in data]
            })
        title = self._generate_title(user_query)
        return {
            "title": {
                "text": title
            },
            "tooltip": {
                "trigger": "axis"
            },
            "legend": {
                "data": y_columns
            },
            "xAxis": {
                "type": "category",
                "data": x_data
            },
            "yAxis": {
                "type": "value"
            },
            "series": series
        }
    
    def _generate_line_chart(self, data: List[Dict[str, Any]], user_query: str) -> Dict[str, Any]:
        bar_config = self._generate_bar_chart(data, user_query)
        for series in bar_config["series"]:
            series["type"] = "line"
        return bar_config
    
    def _generate_pie_chart(self, data: List[Dict[str, Any]], user_query: str) -> Dict[str, Any]:
        features = self._analyze_data_features(data)
        name_column = features["categorical_columns"][0] if features["categorical_columns"] else features["columns"][0]
        value_column = features["numeric_columns"][0] if features["numeric_columns"] else features["columns"][1] if len(features["columns"]) > 1 else features["columns"][0]
        pie_data = [
            {"name": str(row[name_column]), "value": row.get(value_column, 0)}
            for row in data if name_column in row
        ]
        title = self._generate_title(user_query)
        return {
            "title": {
                "text": title,
                "textAlign": "center",
                "left": "center"
            },
            "tooltip": {
                "trigger": "item",
                "formatter": "{a} <br/>{b}: {c} ({d}%)"
            },
            "legend": {
                "orient": "vertical",
                "data": [item["name"] for item in pie_data],
                "right": "right",
                "top": "center"
            },
            "series": [
                {
                    "name": value_column,
                    "type": "pie",
                    "radius": "55%",
                    "center": ["50%", "60%"],
                    "data": pie_data,
                    "emphasis": {
                        "itemStyle": {
                            "shadowBlur": 10,
                            "shadowOffsetX": 0,
                            "shadowColor": "rgba(0, 0, 0, 0.5)"
                        }
                    }
                }
            ]
        }
    
    def _generate_title(self, user_query: str) -> str:
        prompt = PromptTemplate.from_template(
            """根据用户查询生成图表标题。
            用户查询：{user_query}
            请生成一个简洁明了的图表标题，不超过20个字符，确保标题准确反映数据主题和内容。
            标题中不要包括任何特殊字符或格式化。
            只返回标题文本，不要提供任何解释或其他内容。
            """
        )
        chain = prompt | self.llm | StrOutputParser()
        title = chain.invoke({"user_query": user_query}).strip()
        return title if title else "Data Visualization"

class StrOutputParser:    
    def parse(self, text):
        return text.strip()
    def __call__(self, text):
        return self.parse(text)